package common;

import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HRegionLocation;
import org.apache.hadoop.hbase.client.Scan;
import org.apache.hadoop.hbase.mapreduce.TableInputFormat;
import org.apache.hadoop.hbase.mapreduce.TableMapReduceUtil;
import org.apache.hadoop.hbase.mapreduce.TableSplit;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.JobContext;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class MyTableInputFormat extends TableInputFormat {
    @Override
    public List<InputSplit> getSplits(JobContext context) throws IOException {
        this.initialize(context);
        Configuration conf = context.getConfiguration();

        // scan rowkey start 开始范围
        String start = conf.get(TableInputFormat.SCAN_ROW_START, "");
        // scan rowkey end 结束范围
        String end = conf.get(TableInputFormat.SCAN_ROW_STOP, "");
        // 预分区起始
        String splitStart = conf.get("hbase.table.split.startKey", "0");
        // 预分区结束
        String splitEnd = conf.get("hbase.table.split.endKey", "");
        // 分区进制
        int rowKeyRadix = conf.getInt("hbase.table.split.radix", 10);
        // 连接符号
        String rowKeyConcat = conf.get("hbase.table.split.concat", "|");
        // 直接读取指定分区
        String regionSplits = conf.get("hbase.table.split.list", "");
        // 默认Scan
        String scanStr = conf.get(TableInputFormat.SCAN, "");
        int numLength = Math.max(splitEnd.length(), 1);
        String preString = "0000000";

        Scan scan = TableMapReduceUtil.convertStringToScan(scanStr);
        if (StringUtils.isNotBlank(regionSplits) || StringUtils.isNotBlank(splitEnd)) {
            String[] repls;
            if (StringUtils.isNotBlank(regionSplits)) {
                repls = regionSplits.trim().split(",", -1);
            } else {
                Integer s = Integer.parseInt(splitStart);
                Integer e = Long.valueOf(splitEnd, rowKeyRadix).intValue() + 1;
                System.err.println("分区范围------------------");
                repls = IntStream
                        .range(s, e)
                        .mapToObj(value -> {
                            String ss;
                            if (rowKeyRadix == 16) {
                                ss = preString + Integer.toHexString(value);
                            } else {
                                ss = String.valueOf(value);
                            }
                            System.err.println(ss);
                            return StringUtils.substring(ss, ss.length() - numLength, ss.length());
                        }).collect(Collectors.toList()).toArray(new String[]{});
            }
            List<InputSplit> splitLists = new ArrayList<>();
            for (String prefix : repls) {
                HRegionLocation location = getRegionLocator().getRegionLocation(Bytes.toBytes(prefix));
                byte[] _splitStart = Bytes.add(Bytes.toBytes(prefix + rowKeyConcat), Bytes.toBytes(start));
                byte[] _splitStop = Bytes.add(Bytes.toBytes(prefix + rowKeyConcat), Bytes.toBytes(end));

                splitLists.add(new TableSplit(getTable().getName(), scan, _splitStart, _splitStop, location.getHostname()));
            }
            System.err.println("时间戳范围------------------");
            System.err.println(splitLists);
            return splitLists;
        }
        return super.getSplits(context);
    }
}